# -*- coding: utf-8 -*-
import os
import sys
import argparse
import csv

from PIL import Image


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("-d", "--dataset", help="Folder with data", required=True)
    parser.add_argument("-t", "--threshold", help="Threshold for white/total percentage. Default is 0.001", required=False, type=float, default=0.001)
    return parser.parse_args()
    
    
def count_nonblack_pil(img):
    """Return the number of pixels in img that are not black.
    img must be a PIL.Image object in mode RGB.
    """
    bbox = img.getbbox()
    if not bbox: return 0
    return sum(img.crop(bbox)
               .point(lambda x: 255 if x else 0)
               .convert("L")
               .point(bool)
               .getdata())


if __name__ == '__main__':
    args = parse_args()
    threshold = args.threshold

    train_file = open('percentage.csv', 'wb')
    test_writer = csv.writer(train_file, delimiter=';')
    
    count = 0
    for root, dirs, files in os.walk(args.dataset):
        filenames = [os.path.join(root, name) for name in files]
        folder_name = os.path.split(root)[-1].strip(' .0123456789')  #Clear digits and leading\trailing spaces
        if not filenames:
            continue            
        for filename in filenames:    
            count += 1
            im = Image.open(filename)
            n_white_pixels = count_nonblack_pil(im)
            n_pixels = im.size[0]*im.size[1]
            percentage = float(n_white_pixels)/n_pixels
            
            if percentage < threshold:
                #Remove file
                #os.remove(filename)
                test_writer.writerow([filename, folder_name, percentage])
            if count % 500 == 0:
                print "%s done" % count
    
    train_file.close()
        